#include "DependencyEncoder2OGrandChild.h"
#include "math.h"
#include "ObjectReader.h"
#include "DependencyInstance.h"
#include <ctime>

DependencyEncoder2OGrandChild::DependencyEncoder2OGrandChild(int maxLength, DependencyPipe *pipe, char *trainforest, Parameters *param) : DependencyEncoder(maxLength, pipe, trainforest, param) {
	probs_grands = vector<vector<vector<double> > >(maxLength);
	probs_gchs = vector<vector<double> >(maxLength);
	for(int i = 0; i < maxLength; i++){
		probs_gchs[i] = vector<double>(maxLength);
		probs_grands[i] = vector<vector<double> >(maxLength);
		for(int j = 0; j < maxLength; j++){
			int nn;
			if(j == i){
				nn = 0;
			}
			else if(j < i){
				nn = i;
			}
			else{
				nn = maxLength;
			}
			probs_grands[i][j] = vector<double>(nn);
		}
	}
}

DependencyEncoder2OGrandChild::~DependencyEncoder2OGrandChild(){}

int DependencyEncoder2OGrandChild::getKey(int g, int s, int t, int dir, int comp, int length){
	int key = g;
	key = key * length + s;
	key = key * length + t;
	key = key * 2 + dir;
	key = key * 2 + comp;
	return key;
}

double DependencyEncoder2OGrandChild::getZ(int length, double *beta){
	int key, key1, key2;
	for(int s = 0; s < length; s++){
		for(int g = 0; g < length; g++){
			key = getKey(g, s, s, 0, 1, length);
			beta[key] = 0.0;
			key = getKey(g, s, s, 1, 1, length);
			beta[key] = 0.0;
		}
	}

	for(int j = 1; j < length; j++){
		for(int s = 0; s < length && s + j < length; s++){
			int t = s + j;

			double prodProb_st = probs[s][t][0];
			double prodProb_ts = probs[s][t][1];

			for(int g = 0; g < length; g++){
				if(g >= s && g <= t){
					continue;
				}

				//init beta
				//incomplete spans
				int key_gst_0 = getKey(g, s, t, 0, 0, length);
				beta[key_gst_0] = 0.0;
				int key_gts_0 = getKey(g, s, t, 1, 0, length);
				beta[key_gts_0] = 0.0;

				//complete spans
				int key_gst_1 = getKey(g, s, t, 0, 1, length);
				beta[key_gst_1] = 0.0;
				int key_gts_1 = getKey(g, s, t, 1, 1, length);
				beta[key_gts_1] = 0.0;

				bool flg_gst_0 = true, flg_gts_0 = true;
				bool flg_gst_1 = true, flg_gts_1 = true;

				for(int r = s; r < t; r++){
					key1 = getKey(g, s, r, 0, 1, length);
					key2 = getKey(s, r + 1, t, 1, 1, length);
					double prodProb_gst = probs_grands[g][s][t] + probs_gchs[g][t] + prodProb_st;
					beta[key_gst_0] = logsumexp(beta[key_gst_0], beta[key1] + beta[key2] + prodProb_gst, flg_gst_0);
					flg_gst_0 = false;

					key1 = getKey(t, s, r, 0, 1, length);
					key2 = getKey(g, r + 1, t, 1, 1, length);
					double prodProb_gts = probs_grands[g][t][s] + probs_gchs[g][s] + prodProb_ts;
					beta[key_gts_0] = logsumexp(beta[key_gts_0], beta[key1] + beta[key2] + prodProb_gts, flg_gts_0);
					flg_gts_0 = false;
				}

				for(int r = s; r <= t; r++){
					if(r != s){
						key1 = getKey(g, s, r, 0, 0, length);
						key2 = getKey(s, r, t, 0, 1, length);
						beta[key_gst_1] = logsumexp(beta[key_gst_1], beta[key1] + beta[key2], flg_gst_1);
						flg_gst_1 = false;
					}
					if(r != t){
						key1 = getKey(g, r, t, 1, 0, length);
						key2 = getKey(t, s, r, 1, 1, length);
						beta[key_gts_1] = logsumexp(beta[key_gts_1], beta[key1] + beta[key2], flg_gts_1);
						flg_gts_1 = false;
					}
				}
			}
		}
	}
	
	int end = length - 1;
	key = getKey(0, 0, 0, 0, 1, length);
	beta[key] = 0.0;
	key = getKey(end, end, end, 1, 1, length);
	beta[key] = 0.0;

	for(int j = 1; j < length; j++){
		int t = j, s = end - j;
		//init
		double prodProb_0t = probs[0][t][0];
		int key_t_0 = getKey(0, 0, t, 0, 0, length);
		beta[key_t_0] = 0.0;
		int key_t_1 = getKey(0, 0, t, 0, 1, length);
		beta[key_t_1] = 0.0;

		bool flg_t_0 = true, flg_t_1 = true;
		for(int r = 0; r < t; r++){
			key1 = getKey(0, 0, r, 0, 1, length);
			key2 = getKey(0, r + 1, t, 1, 1, length);
			beta[key_t_0] = logsumexp(beta[key_t_0], beta[key1] + beta[key2] + prodProb_0t, flg_t_0);
			flg_t_0 = false;
		}

		for(int r = 1; r <= t; r++){
			key1 = getKey(0, 0, r, 0, 0, length);
			key2 = getKey(0, r, t, 0 ,1, length);
			beta[key_t_1] = logsumexp(beta[key_t_1], beta[key1] + beta[key2], flg_t_1);
			flg_t_1 = false;
		}

		//init
		double prodProb_ns = probs[s][length - 1][1];
		int key_s_0 = getKey(s, s, end, 1, 0, length);
		beta[key_s_0] = 0.0;
		int key_s_1 = getKey(s, s, end, 1, 1, length);
		beta[key_s_1] = 0.0;

		bool flg_s_0 = true, flg_s_1 = true;
		for(int r = s; r < end; r++){
			key1 = getKey(end, s, r, 0, 1, length);
			key2 = getKey(r + 1, r + 1, end, 1, 1, length);
			beta[key_s_0] = logsumexp(beta[key_s_0], beta[key1] + beta[key2] + prodProb_ns, flg_s_0);
			flg_s_0 = false;
		}

		for(int r = s; r < end; r++){
			key1 = getKey(end, s, r, 1, 1, length);
			key2 = getKey(r, r, end, 1, 0, length);
			beta[key_s_1] = logsumexp(beta[key_s_1], beta[key1] + beta[key2], flg_s_1);
			flg_s_1 = false;
		}
	}

	key1 = getKey(0, 0, end, 0, 1, length);
	key2 = getKey(0, 0, end, 1, 1, length);
	return  logsumexp(beta[key1], beta[key2], false);
}

void DependencyEncoder2OGrandChild::inside_outside(int length, double *beta, double *alpha){
	int end = length - 1;
	int key;
	for(int c = 0; c < 2; c++){
		key = getKey(0, 0, end, 0, c, length);
		alpha[key] = 0.0;
		key = getKey(0, 0, end, 1, c, length);
		alpha[key] = 0.0;
	}

	for(int j = end; j >= 1; j--){
		int t = j, s = end - j;

		//init
		int key_t_1 = getKey(0, 0, t, 0, 1, length);
		alpha[key_t_1] = 0.0;
		int key_t_0 = getKey(0, 0, t, 0, 0, length);
		alpha[key_t_0] = 0.0;

		bool flg_t_1 = true, flg_t_0 = true;

		for(int r = t + 1; r < length; r++){
			double prodProb_00r = probs[0][r][0];
			int key_b = getKey(0, t + 1, r, 1, 1, length);
			int key_a = getKey(0, 0, r, 0, 0, length);
			alpha[key_t_1] = logsumexp(alpha[key_t_1], beta[key_b] + alpha[key_a] + prodProb_00r, flg_t_1);
			flg_t_1 = false;
		}

		for(int r = t; r < length; r++){
			int key_b = getKey(0, t, r, 0, 1, length);
			int key_a = getKey(0, 0, r, 0, 1, length);
			alpha[key_t_0] = logsumexp(alpha[key_t_0], beta[key_b] + alpha[key_a], flg_t_0);
			flg_t_0 = false;
		}

		//init
		int key_s_1 = getKey(s, s, end, 1, 1, length);
		alpha[key_s_1] = 0.0;
		int key_s_0 = getKey(s, s, end, 1, 0, length);
		alpha[key_s_0] = 0.0;

		bool flg_s_1 = true, flg_s_0 = true;

		for(int r = 0; r < s; r++){
			double prodProb_nnr = probs[r][end][1];
			int key_b = getKey(end, r, s - 1, 0, 1, length);
			int key_a = getKey(r, r, end, 1, 0, length);
			alpha[key_s_1] = logsumexp(alpha[key_s_1], beta[key_b] + alpha[key_a] + prodProb_nnr, flg_s_1);
			flg_s_1 = false;
		}

		for(int r = 0; r <= s; r++){
			int key_b = getKey(end, r, s, 1, 1, length);
			int key_a = getKey(r, r, end, 1, 1, length);
			alpha[key_s_0] = logsumexp(alpha[key_s_0], beta[key_b] + alpha[key_a], flg_s_0);
			flg_s_0 = false;
		}
	}

	for(int j = end; j >= 1; j--){
		for(int s = 0; s + j < length; s++){
			int t = s + j;
			int key_a, key_b;

			for(int g = 0; g < s; g++){
				//init alpha
				int key_gst_0_1 = getKey(g, s, t, 0, 1, length);
				alpha[key_gst_0_1] = 0.0;
				int key_gst_1_1 = getKey(g, s, t, 1, 1, length);
				alpha[key_gst_1_1] = 0.0;

				bool flg_gst_0_1 = true, flg_gst_1_1 = true;

				//complete spans
				for(int r = 0; r < length; r++){
					if(r > t && r < length){
						double prodProb_gsr = probs_grands[g][s][r] + probs_gchs[g][r] + probs[s][r][0];
						key_b = getKey(s, t + 1, r, 1, 1, length);
						key_a = getKey(g, s, r, 0, 0, length);
						alpha[key_gst_0_1] = logsumexp(alpha[key_gst_0_1], beta[key_b] + alpha[key_a] + prodProb_gsr, flg_gst_0_1);
						flg_gst_0_1 = false;
					}
					if(r > g && r < s){
						double prodProb_gtr = probs_grands[g][t][r] + probs_gchs[g][r] + probs[r][t][1];
						key_b = getKey(t, r, s - 1, 0, 1, length);
						key_a = getKey(g, r, t, 1, 0, length);
						alpha[key_gst_1_1] = logsumexp(alpha[key_gst_1_1], beta[key_b] + alpha[key_a] + prodProb_gtr, flg_gst_1_1);
						flg_gst_1_1 = false;
					}
					if(r < g || r > t){
						key_b = getKey(r, g, s, 0, 0, length);
						key_a = getKey(r, g, t, 0, 1, length);
						alpha[key_gst_0_1] = logsumexp(alpha[key_gst_0_1], beta[key_b] + alpha[key_a], flg_gst_0_1);
						flg_gst_0_1 = false;

						double prodProb_rgt = probs_grands[r][g][t] + probs_gchs[r][t] + probs[g][t][0];
						key_b = getKey(r, g, s - 1, 0, 1, length);
						key_a = getKey(r, g, t, 0, 0, length);
						alpha[key_gst_1_1] = logsumexp(alpha[key_gst_1_1], beta[key_b] + alpha[key_a] + prodProb_rgt, flg_gst_1_1);
						flg_gst_1_1 = false;
					}
				}
				if(g == 0){
					key_b = getKey(0, 0, s, 0, 0, length);
					key_a = getKey(0, 0, t, 0, 1, length);
					alpha[key_gst_0_1] = logsumexp(alpha[key_gst_0_1], beta[key_b] + alpha[key_a], flg_gst_0_1);
					flg_gst_0_1 = false;

					double prodProb_00t = probs[0][t][0];
					key_b = getKey(0, 0, s - 1, 0, 1, length);
					key_a = getKey(0, 0, t, 0, 0, length);
					alpha[key_gst_1_1] = logsumexp(alpha[key_gst_1_1], beta[key_b] + alpha[key_a] + prodProb_00t, flg_gst_1_1);
					flg_gst_1_1 = false;
				}

				//incompleted spans
				int key_gst_0_0 = getKey(g, s, t, 0, 0, length);
				alpha[key_gst_0_0] = 0.0;
				int key_gst_1_0 = getKey(g, s, t, 1, 0, length);
				alpha[key_gst_1_0] = 0.0;

				bool flg_gst_0_0 = true, flg_gst_1_0 = true;
				for(int r = t; r < length; r++){
					key_b = getKey(s, t, r, 0, 1, length);
					key_a = getKey(g, s, r, 0, 1, length);
					alpha[key_gst_0_0] = logsumexp(alpha[key_gst_0_0], beta[key_b] + alpha[key_a], flg_gst_0_0);
					flg_gst_0_0 = false;
				}
				for(int r = g + 1; r <= s; r++){
					key_b = getKey(t, r, s, 1, 1, length);
					key_a = getKey(g, r, t, 1, 1, length);
					alpha[key_gst_1_0] = logsumexp(alpha[key_gst_1_0], beta[key_b] + alpha[key_a], flg_gst_1_0);
					flg_gst_1_0 = false;
				}
			}

			for(int g = t + 1; g < length; g++){
				//init alpha
				int key_gst_0_1 = getKey(g, s, t, 0, 1, length);
				alpha[key_gst_0_1] = 0.0;
				int key_gst_1_1 = getKey(g, s, t, 1, 1, length);
				alpha[key_gst_1_1] = 0.0;

				bool flg_gst_0_1 = true, flg_gst_1_1 = true;

				//completed spans
				for(int r = 0; r < length; r++){
					if(r > t && r < g){
						double prodProb_gsr = probs_grands[g][s][r] + probs_gchs[g][r] + probs[s][r][0];
						key_b = getKey(s, t + 1, r, 1, 1, length);
						key_a = getKey(g, s, r, 0, 0, length);
						alpha[key_gst_0_1] = logsumexp(alpha[key_gst_0_1], beta[key_b] + alpha[key_a] + prodProb_gsr, flg_gst_0_1);
						flg_gst_0_1 = false;
					}
					if(r < s){
						double prodProb_gtr = probs_grands[g][t][r] + probs_gchs[g][r] + probs[r][t][1];
						key_b = getKey(t, r, s - 1, 0, 1, length);
						key_a = getKey(g, r, t, 1, 0, length);
						alpha[key_gst_1_1] = logsumexp(alpha[key_gst_1_1], beta[key_b] + alpha[key_a] + prodProb_gtr, flg_gst_1_1);
						flg_gst_1_1 = false;
					}
					if(r < s || r > g){
						double prodProb_rgs = probs_grands[r][g][s] + probs_gchs[r][s] + probs[s][g][1];
						key_b = getKey(r, t + 1, g, 1, 1, length);
						key_a = getKey(r, s, g, 1, 0, length);
						alpha[key_gst_0_1] = logsumexp(alpha[key_gst_0_1], beta[key_b] + alpha[key_a] + prodProb_rgs, flg_gst_0_1);
						flg_gst_0_1 = false;

						key_b = getKey(r, t, g, 1, 0, length);
						key_a = getKey(r, s, g, 1, 1, length);
						alpha[key_gst_1_1] = logsumexp(alpha[key_gst_1_1], beta[key_b] + alpha[key_a], flg_gst_1_1);
						flg_gst_1_1 = false;
					}
				}
				if(g == end){
					double prodProb_nns = probs[s][end][1];
					key_b = getKey(t + 1, t + 1, end, 1, 1, length);
					key_a = getKey(s, s, end, 1, 0, length);
					alpha[key_gst_0_1] = logsumexp(alpha[key_gst_0_1], beta[key_b] + alpha[key_a] + prodProb_nns, flg_gst_0_1);
					flg_gst_0_1 = false;

					key_b = getKey(t, t, end, 1, 0, length);
					key_a = getKey(s, s, end, 1, 1, length);
					alpha[key_gst_1_1] = logsumexp(alpha[key_gst_1_1], beta[key_b] + alpha[key_a], flg_gst_1_1);
					flg_gst_1_1 = false;
				}

				//incompleted spans
				int key_gst_0_0 = getKey(g, s, t, 0, 0, length);
				alpha[key_gst_0_0] = 0.0;
				int key_gst_1_0 = getKey(g, s, t, 1, 0, length);
				alpha[key_gst_1_0] = 0.0;

				bool flg_gst_0_0 = true, flg_gst_1_0 = true;
				for(int r = t; r < g; r++){
					key_b = getKey(s, t, r, 0, 1, length);
					key_a = getKey(g, s, r, 0 ,1, length);
					alpha[key_gst_0_0] = logsumexp(alpha[key_gst_0_0], beta[key_b] + alpha[key_a], flg_gst_0_0);
					flg_gst_0_0 = false;
				}

				for(int r = 0; r <= s; r++){
					key_b = getKey(t, r, s, 1, 1, length);
					key_a = getKey(g, r, t, 1, 1, length);
					alpha[key_gst_1_0] = logsumexp(alpha[key_gst_1_0], beta[key_b] + alpha[key_a], flg_gst_1_0);
					flg_gst_1_0 = false;
				}
			}
		}
	}
}

void DependencyEncoder2OGrandChild::getExpected(long double *expected, double *beta, double *alpha, double z, int length){
	// Get production crap.
	for (int w1 = 0; w1 < length; w1++) {
		for (int w2 = w1 + 1; w2 < length; w2++) {
			for (int ph = 0; ph < 2; ph++) {
				int *keys = NULL;
				int size = in->readArray(keys);
				double m = 0.0;
				int key;
				for(int g = 0; g < length; g++){
					if(g < w1 || g > w2){
						key = getKey(g, w1, w2, ph, 0, length);
						m += exp(beta[key] + alpha[key] - z);
					}
				}
				if(w1 == 0 && ph == 0){
					key = getKey(0, 0, w2, 0, 0, length);
					m += exp(beta[key] + alpha[key] - z);
				}
				if(w2 == length - 1 && ph == 1){
					key = getKey(w1, w1, length - 1, 1, 0, length);
					m += exp(beta[key] + alpha[key] - z);
				}
				updataExpected(expected, keys, m, size);
				delete[](keys);
			}
		}
	}
	int last = in->readInt();
	if(last != -3){
		cout<<"Error reading file."<<endl;
		exit(0);
	}

	for(int w1 = 0; w1 < length; w1++){
		for(int w2 = w1 + 1; w2 < length; w2++){
			for(int w3 = w1 + 1; w3 < length; w3++){
				if(w3 == w2){
					continue;
				}

				int *keys = NULL;
				int size = in->readArray(keys);
				int key;
				if(w2 < w3){
					key = getKey(w1, w2, w3, 0, 0, length);
				}
				else{
					key = getKey(w1, w3, w2, 1, 0, length);
				}
				double m = exp(beta[key] + alpha[key] - z);
				updataExpected(expected, keys, m, size);
				delete[](keys);
			}
		}

		for(int w2 = w1 - 1; w2 >= 0; w2--){
			for(int w3 = w1 - 1; w3 >= 0; w3--){
				if(w2 == w3){
					continue;
				}

				int *keys = NULL;
				int size = in->readArray(keys);
				int key;
				if(w2 < w3){
					key = getKey(w1, w2, w3, 0, 0, length);
				}
				else{
					key = getKey(w1, w3, w2, 1, 0, length);
				}
				double m = exp(beta[key] + alpha[key] - z);
				updataExpected(expected, keys, m, size);
				delete[](keys);
			}
		}
	}
	last = in->readInt();
	if(last != -3){
		cout<<"Error reading file."<<endl;
		exit(0);
	}

	for(int w1 = 0; w1 < length; w1++){
		for(int w2 = w1 + 1; w2 < length; w2++){
			int *keys = NULL;
			int size = in->readArray(keys);
			double m = 0.0;
			int key;
			for(int w3 = w1 + 1; w3 < length; w3++){
				if(w3 == w2){
					continue;
				}

				if(w3 < w2){
					key = getKey(w1, w3, w2, 0, 0, length);
				}
				else{
					key = getKey(w1, w2, w3, 1, 0, length);
				}
				m += exp(beta[key] + alpha[key] - z);
			}
			updataExpected(expected, keys, m, size);
			delete[](keys);
		}

		for(int w2 = w1 - 1; w2 >= 0; w2--){
			int *keys = NULL;
			int size = in->readArray(keys);
			double m = 0.0;
			int key;
			for(int w3 = w1 - 1; w3 >= 0; w3--){
				if(w2 == w3){
					continue;
				}

				if(w3 < w2){
					key = getKey(w1, w3, w2, 0, 0, length);
				}
				else{
					key = getKey(w1, w2, w3, 1, 0, length);
				}
				m += exp(beta[key] + alpha[key] - z);
			}
			updataExpected(expected, keys, m, size);
			delete[](keys);
		}
	}
	last = in->readInt();
	if(last != -3){
		cout<<"Error reading file."<<endl;
		exit(0);
	}

	int *keys = NULL;
	int size = in->readArray(keys);
	for(int i = 0; i < size; i++){
		int j = keys[i];
		expected[j] -= 1.0;
	}
	
	delete[](keys);
	last = in->readInt();
	if (last != -4) {
		cout<<"Error reading file."<<endl;
		exit(0);
	}

	DependencyInstance* inst = new DependencyInstance();
	inst->readObject(*in);

	last = in->readInt();
	if (last != -1) {
		cout<<"Error reading file."<<endl;
		exit(0);
	}
	delete(inst);
}

double DependencyEncoder2OGrandChild::gradient(long double *expected, int length){
	long offset = in->objTell();
	DependencyInstance *inst = ((DependencyPipe2OGrandChildren*)(pipe))->readInstance(*in, length, probs, probs_grands, probs_gchs, nt_probs, param);
	if(0 != (in->objSeek(offset))){
		cout<<"error on fseek of trainforset!"<<endl;
		exit(0);
	}

	double obj = 0.0;
	double *beta = new double[length * length * length * 2 * 2];
	double *alpha = new double[length * length * length * 2 * 2];

	double z = getZ(length, beta);

	inside_outside(length, beta, alpha);

	obj = z - param->getScore(inst->fv);
	delete(inst);

	getExpected(expected, beta, alpha, z, length);

	delete[](beta);
	delete[](alpha);

	return obj;
}
